import cv2 as cv
import numpy as np


# 内容：稀疏光流估计

cap = cv.VideoCapture("/Users/apple/Desktop/data/vtest.avi")
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)


# ShiTomasi 角点检测参数
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=55,
                      blockSize=7)

# Lucas-Kanade 光流法参数
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))


# 第一帧关键点、创建蒙版画轨迹、特征点随机颜色
prevPts = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_params)
mask = np.zeros_like(prev)
color = np.random.randint(0, 255, (100, 3))


# 遍历视频每一帧
while ret:

    ret, next = cap.read()
    nextGray = cv.cvtColor(next, cv.COLOR_BGR2GRAY)

    # 关键点数量过少，重新寻找
    if (len(prevPts) < 10):
        nextPts = cv.goodFeaturesToTrack(nextGray, mask=None, **feature_params)

    # 计算光流以获取点的新位置
    nextPts, status, err = cv.calcOpticalFlowPyrLK(prevGray, nextGray, prevPts, None, **lk_params)

    # 选择good points作为跟踪点
    prevGPts = prevPts[status == 1]
    nextGPts = nextPts[status == 1]


    # 绘制跟踪曲线
    for i, (new, old) in enumerate(zip(nextGPts, prevGPts)):
        a, b = new.ravel()
        c, d = old.ravel()
        mask = cv.line(mask, (a, b), (c, d), color[i].tolist(), 2)
        next = cv.circle(next, (a, b), 5, color[i].tolist(), -1)

    # 显示
    frame = cv.add(next, mask)
    cv.imshow("1 - Video Detection", frame)

    # 更新上一帧的图像和追踪点
    prevGray = nextGray.copy()
    prevPts = nextGPts.reshape(-1, 1, 2)

    # 按esc退出
    if cv.waitKey(10) & 0xff == 27:
        break


cap.release()
cv.destroyAllWindows()
